from flask import Flask, request, jsonify, Response
from flask_limiter import Limiter
from flask_limiter.util import get_remote_address
import logging
import requests
import time
import schedule
import threading

app = Flask(__name__)

# 创建 robots.txt 的内容
ROBOTS_TXT = """
User-agent: *
Disallow: /
"""

limiter = Limiter(
    app=app,
    key_func=get_remote_address,
    default_limits=["20 per day", "2 per hour", "1 per minute", "1 per 10 second"],
    # storage_uri="redis://47.108.91.31:6379"
)

limiter.init_app(app)

# IP 请求次数记录
ip_request_counts = {}

# 请求时间
last_request_times = {}

# IP 黑名单
blacklisted_ips = set()

# 禁止网段
banned_subnets = ["221.231.47."]

# IP 白名单
whitelisted_ips = set()


# 从文本文件中加载黑名单 IP
def load_blacklisted_ips():
    with open("blacklist.txt", "r") as file:
        blacklisted_ips = set(file.read().splitlines())
    return blacklisted_ips


# 从文本文件中加载白名单 IP
def load_whitelisted_ips():
    with open("whitelist.txt", "r") as file:
        whitelisted_ips = set(file.read().splitlines())
    return whitelisted_ips


# 保存黑名单 IP 到文本文件
def save_blacklisted_ips(blacklisted_ips):
    with open("blacklist.txt", "w") as file:
        file.write("\n".join(blacklisted_ips))


# 保存白名单 IP 到文本文件
def save_whitelisted_ips(whitelisted_ips):
    with open("whitelist.txt", "w") as file:
        file.write("\n".join(whitelisted_ips))


# 添加ip请求次数到： ip_request_counts
def add_request_ips(ip):
    # 更新 IP 请求次数记录
    if ip in ip_request_counts:
        ip_request_counts[ip] += 1
    else:
        ip_request_counts[ip] = 1


# 函数：检测连续增长请求次数的 IP 地址并加入黑名单
def check_and_blacklist_consecutive_ips(ip_request_counts, blacklisted_ips):
    consecutive_count_threshold = 5  # 设定连续增长请求次数的阈值

    consecutive_ips = []
    for ip, count in ip_request_counts.items():
        if count >= consecutive_count_threshold:
            consecutive_ips.append(ip)

    for ip in consecutive_ips:
        blacklisted_ips.add(ip)
        save_blacklisted_ips(blacklisted_ips)
        ip_request_counts.pop(ip, None)
        print(f"IP {ip} has been blacklisted due to consecutive requests")


# 检测IP地址区域
def is_foreign_ip(ip):
    url = f"http://ip-api.com/json/{ip}?fields=status,message,country,countryCode,region,regionName,city,zip,lat,lon,timezone,isp,org,as,proxy"  # 免费接口，不可用时，需找到替代品
    response = requests.get(url)
    data = response.json()

    if 'country' in data:
        print('ip:', ip, '所在区域:', data['country'], '是否为代理:', data['proxy'])
        if data['proxy'] == False:
            if data['countryCode'] != 'CN':
                # 更新 IP 请求次数记录
                # add_request_ips(ip)
                print('ip_request_counts:', ip_request_counts)
                return index_no_limit(ip)
    else:
        return False


# # 每隔30分钟执行一次删除函数,删除 IP 请求次数为 1 的函数
def remove_single_request_ips():
    for ip, count in list(ip_request_counts.items()):
        if count == 1:
            print('到期删除：', ip)
            ip_request_counts.pop(ip, None)


def scheduled_job():
    remove_single_request_ips()


schedule.every(30).minutes.do(scheduled_job)


def scheduled_task():
    while True:
        schedule.run_pending()
        time.sleep(1)


# 启动定时任务线程
remove_single_request_ip_thread = threading.Thread(target=scheduled_task, daemon=True)
remove_single_request_ip_thread.start()

# 初始化黑名单 IP
blacklisted_ips = load_blacklisted_ips()
whitelisted_ips = load_whitelisted_ips()


# 自定义429响应
@app.errorhandler(429)
def handle_rate_limit_exceeded(e):
    return jsonify({"error": "Rate limit exceeded", "message": str(e.description)}), 403



# 路由：处理请求
@app.route('/request_limit', methods=['POST', 'GET'])
# @limiter.limit("1 per 10 second")
def index():
    ip = request.remote_addr
    print('ip:', request.headers['x-forwarded-for'])
    # 校验ip的真实性
    for subnet in banned_subnets:
        if ip.startswith(subnet):
            return jsonify({"message": "this {ip} is blacklisted"}), 403

    # 检测连续增长请求次数的 IP 地址并加入黑名单
    check_and_blacklist_consecutive_ips(ip_request_counts, blacklisted_ips)

    if ip in blacklisted_ips:
        return jsonify({"message": "Your IP is blacklisted"}), 403

    if ip in whitelisted_ips:
        return index_no_limit(ip)

    if is_foreign_ip(ip):
        return jsonify({"message": "Your IP is overseas or No proxy"}), 200

    # 对非白名单 IP 应用速率限制
    limiter.limit("1 per 10 second")(index)

    # 更新 IP 请求次数记录
    add_request_ips(ip)

    print('ip_request_counts:', ip_request_counts)
    print('.......................................')
    print(f"IP: {ip}, Requests: {ip_request_counts[ip]}")
    if request.method == 'GET':
        print("GET请求")
        arges = request.form
        print("参数")
        logging.info(arges)
        print(request)
        print("请求头")
        headers = request.headers
        print(headers)
        logging.info(headers)
        print("请求数据")
        logging.info(request.data)
    elif request.method == 'POST':
        print("POST请求")
        arges = request.form
        print("参数")
        logging.info(arges)
        print(request)
        print("请求头")
        headers = request.headers
        print(headers)
        logging.info(headers)
        print("请求数据")
        logging.info(request.data)
    else:
        print("HEAD请求")
        arges = request.form
        print("参数")
        logging.info(arges)
        print(request)
        print("请求头")
        headers = request.headers
        print(headers)
        logging.info(headers)
        print("请求数据")
        logging.info(request.data)

    if ip_request_counts[ip] >= 10:
        blacklisted_ips.add(ip)
        save_blacklisted_ips(blacklisted_ips)
        ip_request_counts.pop(ip, None)
        return jsonify({"message": f"IP {ip} has reached the limit and has been blacklisted"}), 403

    return jsonify({"message": "Hello, World!"})


# 取消时间限制
@limiter.limit("2 per 1 second")
def index_no_limit(ip):
    if ip in whitelisted_ips:
        return jsonify({"message": f"this {ip} is whitelisted!"}), 200
    return jsonify({"message": "Hello, World!"}), 200


# 添加 IP 到黑名单
@app.route('/blacklist/<ip>', methods=['POST'])
def blacklist_ip(ip):
    blacklisted_ips.add(ip)
    save_blacklisted_ips(blacklisted_ips)
    return jsonify({"message": f"IP {ip} has been blacklisted"})


# 地址 ping 通过
@app.route("/ping")
@limiter.exempt
def ping():
    return "PONG"


# 路由：返回 robots.txt 的内容
@app.route('/robots.txt')
@limiter.limit("1 per 3600 second")
def robots_txt():
    return Response(ROBOTS_TXT, mimetype='text/plain')


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=5000, debug=True)
